
import yaml
from omegaconf import OmegaConf
from omegaconf.dictconfig import DictConfig
from pathlib import Path

from argparse import ArgumentParser


def prepare_args():
    parser=ArgumentParser()
    parser.add_argument(
        "--base_path",
        type=str,
    )
    args=parser.parse_args()
    return args


if __name__ == "__main__":
    args=prepare_args()
    base_cfg=OmegaConf.load(args.base_path)
    
    res=base_cfg.res
    
    del base_cfg.res
    del base_cfg.data
    del base_cfg.lightning

    del base_cfg.model.base_learning_rate
    model_params = base_cfg.model.params
    del model_params.monitor
    try:
        del model_params.ckpt_path
    except:
        pass
    model_params.ddconfig.resolution=res
    
    lossconfig=base_cfg.model.params.lossconfig
    lossconfig.target="torch.nn.Identity"
    del lossconfig.params
    
    save_path = Path(args.base_path).parent / f'inference.yaml'
    OmegaConf.save(base_cfg, str(save_path))